import torch
import torch.nn as nn
import torch.nn.functional as F

class Classfier(nn.Module):
    def __init__(self, input_channel=1024):
        super(Classfier, self).__init__()
        self.softmax = nn.Softmax(dim = 1)
        self.linear_1 = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features=input_channel, out_features=1024),
            nn.ELU(inplace=True)
        )
        self.linear_2 = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features=1024, out_features=256),
            nn.ELU(inplace=True)
        )
        self.linear_3 = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(in_features = 256, out_features=6)
        )
    def forward(self, x):
        x = self.linear_1(x)
        x = self.linear_2(x)
        x = self.linear_3(x)
        x = x.reshape(x.size(0), 1*1*6)
        return x